from langchain_core.output_parsers import JsonOutputParser
from langchain.chat_models import init_chat_model
import os
import asyncio

key = os.getenv("OPENAI_API_KEY")
# print(key)
api_key = str(key)

model = init_chat_model(
    model="gpt-4o-mini",
    model_provider="openai",
    base_url="https://api.zetatechs.com/v1",
    api_key=api_key
)
# Function that does not support streaming.
# It operates on the finalizes inputs rather than
# operating on the input stream.
def _extract_country_names(inputs):
    """A function that does not operates on input streams and breaks streaming."""
    if not isinstance(inputs, dict):
        return ""

    if "countries" not in inputs:
        return ""

    countries = inputs["countries"]

    if not isinstance(countries, list):
        return ""

    country_names = [
        country.get("name") for country in countries if isinstance(country, dict)
    ]
    return country_names


chain = (
    model | JsonOutputParser() | _extract_country_names
)  # This parser only works with OpenAI right now
async def use1():
    async for chunk in chain.astream(
        "output a list of the countries france, spain and japan and their populations in JSON format. "
        'Use a dict with an outer key of "countries" which contains a list of countries. '
        "Each country should have the key `name` and `population`",
    ):
        print(chunk, flush=True)

asyncio.run(use1())


async def use2():
    num_events = 0

    async for event in chain.astream_events(
            "output a list of the countries france, spain and japan and their populations in JSON format. "
            'Use a dict with an outer key of "countries" which contains a list of countries. '
            "Each country should have the key `name` and `population`",
    ):
        kind = event["event"]
        if kind == "on_chat_model_stream":
            print(
                f"Chat model chunk: {repr(event['data']['chunk'].content)}",
                flush=True,
            )
        if kind == "on_parser_stream":
            print(f"Parser chunk: {event['data']['chunk']}", flush=True)
        num_events += 1
        if num_events > 50:
            # Truncate the output
            print("...")
            break

asyncio.run(use2())